#pragma once
#include <utility>
#include <cassert>
namespace wyd
{

    enum Color
    {
        BLACK,
        RED,
    };

    template<class K, class V>
    class RBTreeNode
    {
        typedef RBTreeNode<K, V> Self;
        
        Self* _left;
        Self* _right;
        Self* _parent;
        Color _color;
        std::pair<K, V> _kv;

        RBTreeNode(const std::pair<K, V>& kv)
            : _left(nullptr)
            , _right(nullptr)
            , _parent(nullptr)
            , _color(RED)
            , _kv(kv)
        {}
    };


    template<class K, class V>
    class RBtree
    {
        typedef RBTreeNode<K, V> Node;

    public:
        bool insert(const std::pair<K, V>& kv);

    private:
        void RotateL(Node* parent);
        void RotateR(Node* parent);
    private:
        Node* _root = nullptr;
    };

    template<class K, class V>
    bool RBtree<K, V>::insert(const std::pair<K, V>& kv)
    {
        if (_root == nullptr)
        {
            _root = new Node(kv);
            _root->_color = BLACK;
            return true;
        }

        //找插入位置
        Node* parent = nullptr;
        Node* cur = _root;
        while (cur != nullptr)
        {
            if (cur->_kv.first < kv.first)
            {
                cur = cur->_left;
            }
            else if (cur->_kv.first > kv.first)
            {
                cur = cur->_right;
            }
            else
            {
                assert(false);
                exit(-1);
            }
        }

        //插入节点
        cur = new Node(kv);
        cur->_parent = parent;
        if (parent->_kv.first < kv.first)
        {
            parent->_left = cur;
        }
        else if (parent->_kv.first > kv.first)
        {
            parent->_right = cur;
        }
        else
        {
            assert(false);
            exit(-1);
        }

        //染色
        while (parent && parent->_color == RED)
        {
            Node* granparent = parent->_parent;
            if (parent == granparent->_left)
            {
                Node* uncle = granparent->_right;
                if (uncle && uncle->_color == RED)
                {
                    parent->_color = uncle->_color = BLACK;
                    if (granparent != _root)
                        granparent->_color = RED;

                    cur = granparent;
                    parent = cur->_parent;
                }
                else // uncle 不存在 或 为黑
                {
                    if (cur == parent->_left) 
                    {
                        RotateR(granparent);
                    
                        parent->_color = BLACK;
                        granparent->_color = RED;
                    }
                    else
                    {
                        RotateL(parent);
                        RotateR(granparent);

                        cur->_color = BLACK;
                        granparent->_color = RED;

                    }
                    break;
                }
            }
            else // parent == granparent->_right
            {
                Node* uncle = granparent->_left;
                if (uncle && uncle->_color == RED)
                {
                    parent->_color = uncle->_color = BLACK;
                    if (granparent != _root)
                        granparent->_color = RED;

                    cur = granparent;
                    parent = cur->_parent;
                }
                else // uncle 不存在 或 为黑
                {
                    if (cur == parent->_right) 
                    {
                        RotateL(granparent);
                    
                        parent->_color = BLACK;
                        granparent->_color = RED;
                    }
                    else
                    {
                        RotateR(parent);
                        RotateL(granparent);

                        cur->_color = BLACK;
                        granparent->_color = RED;
                    }
                    break;
                }
            }
        }

    }

    template<class K, class V>
    void RBtree<K, V>::RotateL(Node* parent)
    {
        Node* son = parent->_right;
        Node* sonL = son->_left;
        
        parent->_right = sonL;
        if (sonL != nullptr)
        {
            sonL->_parent = parent;
        }
        
        Node* granparent = parent->_parent;
        son->_parent = granparent;
        if (granparent != nullptr)
        {
            if (granparent->_left == parent)
            {
                granparent->_left = son;
            }
            else if (granparent->_right == parent)
            {
                granparent->_right = son;
            }
        }
        else
        {
            _root == son;
        }

        son->_left = parent;
        parent->_parent = son;
    }

    template<class K, class V>
    void RBtree<K, V>::RotateR(Node* parent)
    {
        Node* son = parent->_left;
        Node* sonR = son->_right;

        parent->_left = sonR;
        if (sonR != nullptr)
        {
            sonR->_parent = parent;
        }
        
        Node* granparent = parent->_parent;
        son->_parent = granparent;
        if (granparent != nullptr)
        {
            if (granparent->_left == parent)
            {
                granparent->_left = son;
            }
            else if (granparent->_right == parent)
            {
                granparent->_right = son;
            }
        }
        else
        {
            _root = son;
        }

        son->_right = parent;
        parent->_parent = son;
    }
}